package cn.virens.web.components.websocket;

import java.security.Principal;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;

import cn.hutool.core.date.DateUtil;

public class WebSocketHandler extends AbstractWebSocketHandler implements WebSocketPool {
	private static final Logger logger = LoggerFactory.getLogger(WebSocketHandler.class);
	private static final Map<String, Map<String, WebSocketSession>> pool = new ConcurrentHashMap<>();

	@Override
	public void afterConnectionEstablished(WebSocketSession session) throws Exception {
		Principal principal = session.getPrincipal();
		if (principal instanceof SimplePrincipal) {
			logger.debug("连接成功:" + session.getRemoteAddress());
			logger.debug("连接用户:" + principal.toString());

			String group = SimplePrincipal.getAccount(principal);
			String account = SimplePrincipal.getAccount(principal);

			synchronized (pool) {
				pool.computeIfAbsent(group, (key) -> {
					return new HashMap<>();
				}).put(account, session);
			}
		} else {
			logger.error("连接无效:" + session.getRemoteAddress());
			logger.error("登录信息无效");

			// 关闭无效链接
			session.close();
		}
	}

	@Override
	protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
		JSONObject payload = JSON.parseObject(message.getPayload());
		logger.debug("收到消息.content:" + payload.getString("content"));
		logger.debug("收到消息.sign:" + payload.getString("sign"));

		Map<String, Object> a = session.getAttributes();
		if (a != null) {
			payload.put("self", true);
			payload.put("success", true);
			payload.put("createTime", DateUtil.now());
			payload.put("name", WebSocketUtil.getName(session));
			payload.put("avatar", WebSocketUtil.getAvatar(session));
			payload.put("account", WebSocketUtil.getAccount(session));
		} else {
			payload.put("success", false);
			payload.put("message", "未登录");
		}

		// 将消息回传到用户
		this.sendMessage(session, payload);
	}

	@Override
	public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
		Principal principal = session.getPrincipal();
		if (principal instanceof SimplePrincipal) {
			logger.debug("连接关闭:" + session.getRemoteAddress());
			logger.debug("断开用户:" + principal.toString());

			String group = ((SimplePrincipal) principal).getGroup();
			String account = ((SimplePrincipal) principal).getAccount();

			synchronized (pool) {
				if (pool.containsKey(group)) {
					Map<String, WebSocketSession> groupMap = pool.remove(group);

					if (groupMap != null && groupMap.containsKey(account)) {
						groupMap.remove(account);
					}
				}
			}
		} else {
			logger.debug("连接已断开");
		}
	}

	@Override
	public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
		logger.error("连接异常", exception);
	}

	@Override
	public int sendAll(JSONObject json) {
		int result = 0;
		synchronized (pool) {
			for (Map<String, WebSocketSession> groupMap : pool.values()) {
				for (WebSocketSession session : groupMap.values()) {
					result += sendMessage(session, json);
				}
			}
		}

		return result;
	}

	@Override
	public int sendGroup(String group, JSONObject json) {
		int result = 0;
		synchronized (pool) {
			Map<String, WebSocketSession> groupMap = pool.get(group);
			// 如果组不存在，直接返回
			if (groupMap == null) { return result; }
			// 对组下面的成员发送信息
			for (WebSocketSession session : groupMap.values()) {
				result += sendMessage(session, json);
			}
		}
		return result;
	}

	@Override
	public int sendAccount(String account, JSONObject json) {
		int result = 0;
		synchronized (pool) {
			for (Map<String, WebSocketSession> groupMap : pool.values()) {
				result += sendMessage(groupMap.get(account), json);
			}
		}
		return result;
	}

	private int sendMessage(WebSocketSession session, JSONObject json) {
		try {
			return WebSocketUtil.sendMessage(session, json);
		} catch (Exception e) {
			logger.error("发送失败：" + e.getMessage());
		}

		return 0;
	}

}